import pandas as pd


class DataAccessor:
    """数据访问器，提供数据访问接口"""

    def __init__(self, regions_df, contracts_df, assets_df, time_periods):
        self.regions_df = regions_df
        self.contracts_df = contracts_df
        self.assets_df = assets_df
        self.time_periods = time_periods

        # 预计算一些常用的数据视图
        self._precompute_views()

    def _precompute_views(self):
        """预计算常用的数据视图，提高查询效率"""
        # 按业务类型分组的合同视图
        self.contracts_by_business = self.contracts_df.groupby('business_type')

        # 按区域分组的合同视图
        self.contracts_by_region = self.contracts_df.groupby('region_id')

        # 按时间周期分组的合同视图
        self.contracts_by_period = self.contracts_df.groupby('period')

        # 按区域和时间周期分组的合同视图
        self.contracts_by_region_period = self.contracts_df.groupby(['region_id', 'period'])

    def get_regions(self):
        """获取所有区域信息"""
        return self.regions_df

    def get_region_id(self, region_name):
        """根据区域名称获取区域ID"""
        return self.regions_df[self.regions_df["region_name"] == region_name]["region_id"].values[0]

    def get_contracts(self, region_id=None, period=None, business_type=None):
        """获取合同数据，可以按区域、周期和业务类型筛选"""
        contracts = self.contracts_df.copy()

        if region_id:
            contracts = contracts[contracts["region_id"] == region_id]
        if period:
            contracts = contracts[contracts["period"] == period]
        if business_type:
            if isinstance(business_type, list):
                contracts = contracts[contracts["business_type"].isin(business_type)]
            else:
                contracts = contracts[contracts["business_type"] == business_type]

        return contracts

    def get_assets(self, region_id=None, period=None):
        """获取资产数据，可以按区域和周期筛选"""
        assets = self.assets_df.copy()

        if region_id:
            assets = assets[assets["region_id"] == region_id]
        if period:
            assets = assets[assets["period"] == period]

        return assets

    def get_time_periods(self):
        """获取所有时间周期"""
        return self.time_periods

    def get_latest_period(self):
        """获取最新的时间周期"""
        return self.time_periods[-1]
